#!/usr/bin/env julia

"""
DRP vs 3DVar synthetic comparison using ensemble-reduced control space.
"""

using LinearAlgebra
using Random
using Statistics
using Printf
using Dates

Random.seed!(20251008)

# Problem dimensions
n_state = parse(Int, get(ENV, "DRP_STATE_SIZE", "2000"))
n_obs_per_window = parse(Int, get(ENV, "DRP_OBS_PER_WINDOW", "400"))
n_windows = parse(Int, get(ENV, "DRP_WINDOWS", "3"))
ensemble_size = parse(Int, get(ENV, "DRP_ENSEMBLE", "40"))
n_obs = n_obs_per_window * n_windows

println("="^72)
println("DRP vs 3DVar Synthetic Comparison")
println("="^72)
println("State size: $n_state, Observations: $n_obs ($n_windows windows)")
println("Ensemble size (DRP): $ensemble_size")

# Construct synthetic truth/background
x_true = randn(n_state) .* 2.0
background_spread = 1.5
x_background = x_true + randn(n_state) .* background_spread

# Background error covariance (diagonal with smooth variation)
variances = @. 1.0 + 0.5 * sin((1:n_state) / 150)
B = Diagonal(variances)
B_inv = Diagonal(1 ./ variances)

# Observation operator: random subsampling with smoothing by window
function build_observation_matrix(n_state, n_obs_per_window, n_windows)
    H = zeros(n_obs_per_window * n_windows, n_state)
    obs_idx = 1
    for window in 1:n_windows
        offset = (window - 1) * floor(Int, n_state / n_windows)
        idxs = randperm(n_state)[1:n_obs_per_window]
        idxs = ((idxs .+ offset .- 1) .% n_state) .+ 1
        for i in idxs
            H[obs_idx, i] = 1.0
            obs_idx += 1
        end
    end
    return H
end

H = build_observation_matrix(n_state, n_obs_per_window, n_windows)

# Observation error covariance
obs_error_std = 1.2
R = Diagonal(fill(obs_error_std^2, n_obs))
R_inv = Diagonal(fill(1 / obs_error_std^2, n_obs))

# Generate observations with noise
noise = randn(n_obs) .* obs_error_std
obs_background = H * x_true + noise
innovation_b = obs_background - H * x_background

function cost_function(x, x_b, B_inv, H, y, R_inv)
    dx = x - x_b
    innov = H * x - y
    return 0.5 * (dx' * (B_inv * dx)) + 0.5 * (innov' * (R_inv * innov))
end

cost_background = cost_function(x_background, x_background, B_inv, H, obs_background, R_inv)

println("Background cost J(x_b): $(round(cost_background, digits=3))")

# 3DVar closed-form analysis
start_time_3dvar = time()
K = B * H' * inv(H * B * H' + R)
analysis_3dvar = x_background + K * innovation_b
execution_time_3dvar = time() - start_time_3dvar
cost_3dvar = cost_function(analysis_3dvar, x_background, B_inv, H, obs_background, R_inv)

# DRP reduced-space solution
println("Computing DRP reduced-space analysis...")
start_time = time()

sqrtB_vec = sqrt.(variances)
sqrtB_mat = Diagonal(sqrtB_vec)
svd_components = svd(H * sqrtB_mat)
V_red = svd_components.V[:, 1:ensemble_size]
ensemble_perturbations = sqrtB_mat * V_red
HP = H * ensemble_perturbations
S = I + HP' * (R_inv * HP)
alpha = S \ (HP' * (R_inv * innovation_b))
analysis_drp = x_background + ensemble_perturbations * alpha
execution_time_drp = time() - start_time
cost_drp = cost_function(analysis_drp, x_background, B_inv, H, obs_background, R_inv)

println(@sprintf("DRP solve time: %.3f s", execution_time_drp))

# Metrics
function innovation_rms(x)
    resid = H * x - obs_background
    return sqrt(mean(resid.^2))
end

metrics = Dict(
    "background" => (
        cost = cost_background,
        rms_innov = innovation_rms(x_background),
        norm_increment = 0.0,
        solve_time = 0.0,
    ),
    "3dvar" => (
        cost = cost_3dvar,
        rms_innov = innovation_rms(analysis_3dvar),
        norm_increment = norm(analysis_3dvar - x_background),
        solve_time = execution_time_3dvar,
    ),
    "drp4dvar" => (
        cost = cost_drp,
        rms_innov = innovation_rms(analysis_drp),
        norm_increment = norm(analysis_drp - x_background),
        solve_time = execution_time_drp,
        reduced_rank = ensemble_size,
    ),
)

println("\nMetrics (cost, RMS innovation, increment norm):")
for (label, data) in metrics
    println(@sprintf("  %-10s cost=%10.3f  rms=%8.4f  |dx|=%8.3f", label, data.cost, data.rms_innov, data.norm_increment))
end

speedup = metrics["3dvar"].norm_increment == 0 ? NaN : metrics["3dvar"].norm_increment / metrics["drp4dvar"].norm_increment
improvement = 100 * (metrics["background"].cost - metrics["drp4dvar"].cost) / metrics["background"].cost

println(@sprintf("\nCost improvement vs background: %.2f%%", improvement))
println(@sprintf("Innovation reduction (background → DRP): %.2f%%", 100 * (metrics["background"].rms_innov - metrics["drp4dvar"].rms_innov) / metrics["background"].rms_innov))
println(@sprintf("DRP vs 3DVar cost ratio: %.3f", metrics["drp4dvar"].cost / metrics["3dvar"].cost))
println(@sprintf("3DVar solve time: %.3f s", metrics["3dvar"].solve_time))
println(@sprintf("DRP solve time: %.3f s", metrics["drp4dvar"].solve_time))

# Export CSV summary
using DelimitedFiles
results_dir = joinpath(@__DIR__, "..", "..", "results", "drp4dvar_vs_3dvar")
mkpath(results_dir)
timestamp = Dates.format(Dates.now(), "yyyymmdd_HHMMSS")

csv_path = joinpath(results_dir, "metrics_$timestamp.csv")
open(csv_path, "w") do io
    println(io, "method,cost,rms_innovation,increment_norm,solve_time,reduced_rank")
    println(io, @sprintf("background,%.6f,%.6f,%.6f,%.6f,%d", metrics["background"].cost, metrics["background"].rms_innov, metrics["background"].norm_increment, metrics["background"].solve_time, 0))
    println(io, @sprintf("3dvar,%.6f,%.6f,%.6f,%.6f,%d", metrics["3dvar"].cost, metrics["3dvar"].rms_innov, metrics["3dvar"].norm_increment, metrics["3dvar"].solve_time, n_state))
    println(io, @sprintf("drp4dvar,%.6f,%.6f,%.6f,%.6f,%d", metrics["drp4dvar"].cost, metrics["drp4dvar"].rms_innov, metrics["drp4dvar"].norm_increment, metrics["drp4dvar"].solve_time, ensemble_size))
end

println("\nSaved metrics to $csv_path")
